#define DOF 3

#include "test.hpp"
#include "stencil_patterns.hpp"

int main(int argc, char* argv[])
{
    int my_pid, num_procs;
    MPI_Init(&argc, &argv);
    MPI_Comm_rank(MPI_COMM_WORLD, &my_pid);
    MPI_Comm_size(MPI_COMM_WORLD, &num_procs);

    int arg_cnt = 1;
    const std::string case_name(argv[arg_cnt++]);
    idx_t num_offset = -1;
    const int * _stencil_offset = nullptr;
    if      (case_name == "DEMO3d7" ) { num_offset =  7; _stencil_offset = stencil_offset_3d7;     }
    else if (case_name == "DEMO3d15") { num_offset = 15; _stencil_offset = stencil_offset_3d15;    }
    else if (case_name == "DEMO3d19") { num_offset = 19; _stencil_offset = stencil_offset_3d19;    }
    else if (case_name == "DEMO3d27") { num_offset = 27; _stencil_offset = stencil_offset_3d27;    }
    else if (case_name == "DEMO2d5" ) { num_offset =  5; _stencil_offset = stencil_offset_2d5i1i2; }
    else MPI_Abort(MPI_COMM_WORLD, -1111);

    // 三维划分
    // 注意这里的三维大小的第0和第1维应当与原来的SysSMG程序调转
    int case_dim    [3] = { atoi(argv[arg_cnt++]), atoi(argv[arg_cnt++]), atoi(argv[arg_cnt++]) };
    int num_proc_dim[3] = { atoi(argv[arg_cnt++]), atoi(argv[arg_cnt++]), atoi(argv[arg_cnt++]) };
    assert(num_proc_dim[0] * num_proc_dim[1] * num_proc_dim[2] == num_procs);

    idx_t cart_ids[3];
    {// 计算本进程在part内处于什么位置
        cart_ids[0] =  my_pid / (num_proc_dim[1] * num_proc_dim[2]);
        cart_ids[1] = (my_pid - cart_ids[0] * num_proc_dim[1] * num_proc_dim[2]) / num_proc_dim[2];
        cart_ids[2] = (my_pid - cart_ids[0] * num_proc_dim[1] * num_proc_dim[2] - cart_ids[1] * num_proc_dim[2]);
    }
    
    idx_t my_ilower[3], my_iupper[3], box_ends[3], my_dims[3];
    idx_t my_nelems = 1;
    for (idx_t d = 0; d < 3; d++) {
        idx_t avg_load = case_dim[d] / num_proc_dim[d];
        idx_t remain   = case_dim[d] - avg_load * num_proc_dim[d];
    
        my_ilower[d] =  cart_ids[d]     * avg_load;
        idx_t my_load = -1;
        if (cart_ids[d] < remain) { my_ilower[d] += cart_ids[d]; my_load = avg_load + 1; } 
        else                      { my_ilower[d] +=   remain   ; my_load = avg_load    ; }

        my_iupper[d] = my_ilower[d] + my_load - 1;
        box_ends[d] = my_iupper[d] + 1;
        my_dims[d] = my_iupper[d] - my_ilower[d] + 1;
        my_nelems *= my_dims[d];
    }
    const idx_t tot_elems = case_dim[0] * case_dim[1] * case_dim[2]; assert(tot_elems > 0);
    idx_t * row_partits = new idx_t [num_procs + 1];
    row_partits[0] = 0;
    constexpr MPI_Datatype mpi_idx_type = sizeof(idx_t) == 8 ? MPI_LONG_LONG : MPI_INT;
    MPI_Allgather(& my_nelems, 1, mpi_idx_type, row_partits + 1, 1, mpi_idx_type, MPI_COMM_WORLD);
    for (int p = 0; p < num_procs; p++)
        row_partits[p + 1] += row_partits[p];
    assert(tot_elems == row_partits[num_procs]);

    idx_t * proc_dims = new idx_t [3 * num_procs];
    idx_t * proc_begs = new idx_t [3 * num_procs];
    MPI_Allgather(my_dims  , 3, mpi_idx_type, proc_dims, 3, mpi_idx_type, MPI_COMM_WORLD);
    MPI_Allgather(my_ilower, 3, mpi_idx_type, proc_begs, 3, mpi_idx_type, MPI_COMM_WORLD);
    assert( proc_dims[my_pid*3 + 0] == my_dims[0] &&
            proc_dims[my_pid*3 + 1] == my_dims[1] &&
            proc_dims[my_pid*3 + 2] == my_dims[2]);
    
    const std::string   its_name = std::string(argv[arg_cnt++]),
                        prc_name = std::string(argv[arg_cnt++]);
    std::string config_mg_file;
	if (prc_name == "GMG") config_mg_file = std::string(argv[arg_cnt++]);
    
    if (my_pid == 0) printf("Max threads: %d\n", omp_get_max_threads());

    // for (int p = 0; p < num_procs; p ++) {
    //     if (my_pid == p) {
    //         printf("Proc %d [%d, %d) x [%d, %d) x [%d, %d)\n", my_pid,
    //             my_ilower[0], my_iupper[0] + 1, 
    //             my_ilower[1], my_iupper[1] + 1,
    //             my_ilower[2], my_iupper[2] + 1 );
    //     }
    //     MPI_Barrier(MPI_COMM_WORLD);
    // }

    {
        par_CSRMatrix<idx_t, ksp_t, ksp_t, DOF> A(MPI_COMM_WORLD,   tot_elems, row_partits[my_pid], row_partits[my_pid + 1],
                                                                    tot_elems, row_partits[my_pid], row_partits[my_pid + 1] );
        par_Vector<idx_t, ksp_t, DOF> x0(MPI_COMM_WORLD,            tot_elems, row_partits[my_pid], row_partits[my_pid + 1] );
        par_Vector<idx_t, ksp_t, DOF> x(x0), b(x0), y(x0);

        x0.set_val(0.0);
        x.set_val(0.0); b.set_val(1.0); y.set_val(0.0);

        {
            long long * dist_rpt = new long long [my_nelems + 1]; dist_rpt[0] = 0;

            for (idx_t loc_1D = 0; loc_1D < my_nelems; loc_1D++) {
                idx_t loc_3D[3];
                loc_3D[0] =  loc_1D / (my_dims[1] * my_dims[2]);
                loc_3D[1] = (loc_1D -  my_dims[1] * my_dims[2] * loc_3D[0]) / my_dims[2];
                loc_3D[2] =  loc_1D -  my_dims[1] * my_dims[2] * loc_3D[0]  - my_dims[2] * loc_3D[1];
                idx_t glb_3D[3] = {
                    loc_3D[0] + my_ilower[0],
                    loc_3D[1] + my_ilower[1],
                    loc_3D[2] + my_ilower[2]
                };
                idx_t ngb_glb_3D[3];
                idx_t cnt = 0;
                for (idx_t k = 0; k < num_offset; k ++) {
                    ngb_glb_3D[0] = glb_3D[0] + _stencil_offset[k * 3 + 0];
                    ngb_glb_3D[1] = glb_3D[1] + _stencil_offset[k * 3 + 1];
                    ngb_glb_3D[2] = glb_3D[2] + _stencil_offset[k * 3 + 2];
                
                    if (0 <= ngb_glb_3D[0] && ngb_glb_3D[0] < case_dim[0] &&
                        0 <= ngb_glb_3D[1] && ngb_glb_3D[1] < case_dim[1] &&
                        0 <= ngb_glb_3D[2] && ngb_glb_3D[2] < case_dim[2])
                    {
                        cnt ++;
                    }
                }
                dist_rpt[loc_1D + 1] = cnt;
            }

            for (idx_t loc_1D = 0; loc_1D < my_nelems; loc_1D++) {
                dist_rpt[loc_1D + 1] += dist_rpt[loc_1D];
            }
            
            constexpr int e_size = DOF*DOF;
            idx_t * dist_cid = new idx_t [dist_rpt[my_nelems]         ];
            ksp_t * dist_val = new ksp_t [dist_rpt[my_nelems] * e_size];

            for (idx_t loc_1D = 0; loc_1D < my_nelems; loc_1D++) {
                const idx_t glb_1D = loc_1D + row_partits[my_pid];
                idx_t loc_3D[3];
                loc_3D[0] =  loc_1D / (my_dims[1] * my_dims[2]);
                loc_3D[1] = (loc_1D -  my_dims[1] * my_dims[2] * loc_3D[0]) / my_dims[2];
                loc_3D[2] =  loc_1D -  my_dims[1] * my_dims[2] * loc_3D[0]  - my_dims[2] * loc_3D[1];
                idx_t glb_3D[3] = {
                    loc_3D[0] + my_ilower[0],
                    loc_3D[1] + my_ilower[1],
                    loc_3D[2] + my_ilower[2]
                };
                idx_t ngb_glb_3D[3], ngb_loc_3D[3];
                idx_t cnt = 0;
                for (idx_t k = 0; k < num_offset; k ++) {
                    ngb_glb_3D[0] = glb_3D[0] + _stencil_offset[k * 3 + 0];
                    ngb_glb_3D[1] = glb_3D[1] + _stencil_offset[k * 3 + 1];
                    ngb_glb_3D[2] = glb_3D[2] + _stencil_offset[k * 3 + 2];
                    if (0 <= ngb_glb_3D[0] && ngb_glb_3D[0] < case_dim[0] &&
                        0 <= ngb_glb_3D[1] && ngb_glb_3D[1] < case_dim[1] &&
                        0 <= ngb_glb_3D[2] && ngb_glb_3D[2] < case_dim[2])
                    {
                        int ngb_proc_cart[3];
                        for (int d = 0; d < 3; d++) {
                            if      (ngb_glb_3D[d] < my_ilower[d])  ngb_proc_cart[d] = cart_ids[d] - 1;
                            else if (ngb_glb_3D[d] > my_iupper[d])  ngb_proc_cart[d] = cart_ids[d] + 1;
                            else                                    ngb_proc_cart[d] = cart_ids[d];
                        }
                        const int ngb_pid = (ngb_proc_cart[0] * num_proc_dim[1] + ngb_proc_cart[1]) * num_proc_dim[2] + ngb_proc_cart[2];
                        const idx_t * ngb_dims = proc_dims + ngb_pid*3;
                        const idx_t * ngb_begs = proc_begs + ngb_pid*3;
                        for (int d = 0; d < 3; d++) {
                            ngb_loc_3D[d] = ngb_glb_3D[d] - ngb_begs[d];
                            assert(0 <= ngb_loc_3D[d] && ngb_loc_3D[d] < ngb_dims[d]);
                        }
                        const idx_t ngb_glb_1D = row_partits[ngb_pid] + 
                            (ngb_loc_3D[0] * ngb_dims[1] + ngb_loc_3D[1]) * ngb_dims[2] + ngb_loc_3D[2];

                        dist_cid[dist_rpt[loc_1D] + cnt] =  ngb_glb_1D;
                        ksp_t v = (ngb_glb_1D == glb_1D) ? (num_offset - 1) : -1.0;
                        ksp_t * dst_ptr = dist_val + (dist_rpt[loc_1D] + cnt) * e_size;
                        dst_ptr[0] = v; dst_ptr[1] = 0; dst_ptr[2] = 0;
                        dst_ptr[3] = 0; dst_ptr[4] = v; dst_ptr[5] = 0;
                        dst_ptr[6] = 0; dst_ptr[7] = 0; dst_ptr[8] = v;

                        cnt ++;
                    }
                }
                assert(cnt + dist_rpt[loc_1D] == dist_rpt[loc_1D + 1]);
            }

            A.set_values_distributed(dist_rpt, dist_cid, dist_val);
            delete [] dist_rpt; delete [] dist_cid; delete [] dist_val;
        }
        

        const int test_cnt = 1;
        std::vector<TEST_RECORD> records;
        for (int test = 0; test < test_cnt; test++) {
            vec_copy(x0, x);

            double ck = vec_dot<idx_t, ksp_t, double>(b, b);
            if (my_pid == 0) printf("(  b,   b) = %.15e\n", ck);

            ck = vec_dot<idx_t, ksp_t, double>(x, x);
            if (my_pid == 0) printf("(  x,   x) = %.15e\n", ck);

            const idx_t __warm_cnt = 3;
            for (int __test = 0; __test < __warm_cnt; __test ++) {
                A.Mult(b, y, false);
            }

            const idx_t __perf_cnt = 10;
            double t_perf_spmv = wall_time();
            for (int __test = 0; __test < __perf_cnt; __test ++) {
                A.Mult(b, y, false);
            }
            t_perf_spmv = wall_time() - t_perf_spmv;
            double mem_size = (
            (double)   A.diag.nrows      * 2 * DOF * sizeof(ksp_t) // 向量
            +(double) (A.diag.nrows + 1) *           sizeof(idx_t) // 对角部分
            +(double) (A.diag.nnz      ) *           sizeof(idx_t)
            +(double) (A.diag.nnz      ) * DOF*DOF * sizeof(ksp_t)
            +(double) (A.offd.nrows + 1) *           sizeof(idx_t) // 非对角部分
            +(double) (A.offd.nnz      ) *           sizeof(idx_t)
            +(double) (A.offd.nnz      ) * DOF*DOF * sizeof(ksp_t)
            ) / (1024*1024*1024);// GB
            double glb_mem_size = 0;
            MPI_Allreduce(& mem_size, & glb_mem_size, 1, MPI_DOUBLE, MPI_SUM, MPI_COMM_WORLD);
            if (my_pid == 0) {
                double t_single = t_perf_spmv / __perf_cnt;
                double effectBW = glb_mem_size / t_single;
                printf("perf_spmv %.5f s per count mem size %.2f GB %.2f GB/s\n", t_single,
                    glb_mem_size, effectBW);
            }
            // Max 9462 * 2 纯多进程（64进程）302.64 GB/s
            
            ck = vec_dot<idx_t, ksp_t, double>(y, y);
            if (my_pid == 0) printf("(A*b, A*b) = %.27e\n", ck);

            A.Mult(x, y, false);
            ck = vec_dot<idx_t, ksp_t, double>(y, y);
            if (my_pid == 0) printf("(A*x, A*x) = %.15e\n", ck);

            TEST_CONFIG config;
            config.config_mg_file = config_mg_file;
            config.max_iter = 30;
            config.rtol = 1.0e-6;
            TEST_RECORD rec;

            buildup_solver(its_name, prc_name, config);
            setup_and_solve(A, b, x, rec);

            double true_r_norm, b_norm;
            check_residual(A, x, b, y, true_r_norm, b_norm);

            if (my_pid == 0) {
                printf("\033[1;35mtrue ||r|| = %20.16e ||r||/||b||= %20.16e\033[0m\n", 
                    true_r_norm, true_r_norm / b_norm);
                printf("Proc %d Setup, Solve costs %.6f %.6f s\n",
                    my_pid, rec.setup, rec.solve);
            }

            // double lambda_A  = check_eigen_value(50, A);
            // double lambda_MA = check_preconditioned_eigen_value(50, A);
            // if (my_pid == 0) {
            //     printf("Max Eigen value = %.10e\n", lambda_A);
            //     printf("Max Eigen value after Precond = %.10e\n", lambda_MA);
            // }

            stat_part_times(rec);
            records.push_back(rec);

            destroy_solver();
        }// test loop

        // 输出最佳时间
        if (my_pid == 0) {
            stat_multi_runs(records);
        }
    }
    MPI_Finalize();
    return 0;
}
